import mindspore
from mindnlp import core
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
import gradio as gr
import os
import time
from threading import Thread

from RAG import chunking, embedding, vectordb, retrive, rerank, generation

def rag_pipeline(
    query: str,
    data_path: str = "./data/data.json",
    db_path: str = "./vector_db",
    top_k_retrive=20,
    top_k_rerank=10,
):
    """
    RAG完整流程函数

    Args:
        query (str): 用户查询
        data_path (str): 数据文件路径
        db_path (str): 向量数据库路径
        top_k_retrive = 6
        top_k_rerank = 4

    Returns:
        str: 生成的参考
    """


    # 如果数据库不存在，则初始化
    if not os.path.exists(db_path) or not os.listdir(db_path):
        chunks = chunking.chunking(data_path)
        data_embedding = embedding.embedding(chunks)
        vectordb.store_to_db(db_path, data_embedding, chunks)

    # 执行RAG流程
    data_retrive = retrive.retrive(query, db_path, top_k_retrive)
    data_rerank = rerank.rerank(query, data_retrive, top_k_rerank)

    return "\n\n".join(data_rerank)

# --- 配置区域 ---
# 1. 模型路径：请确保该路径下包含模型文件和配置文件
MODEL_PATH = './DeepSeek-R1-Distill-Qwen-1.5B'
# 2. System Prompt：定义模型的角色和行为，可以根据需要修改
SYSTEM_PROMPT = "按照用户的要求进行回答"
# 3. Gradio 界面文本：可以根据需要修改界面上的介绍和样例
GRADIO_TITLE = "🎓 XDUMsgBot - 西安电子科技大学问答助手"
GRADIO_EXAMPLES = [
                "学校的奖学金政策是什么？",
                "如何申请研究生？",
                "本科生的培养方案有哪些？",
                "学校有哪些特色专业？",
                "校园生活怎么样？",
            ]

# --- 1. 模型一次性加载 ---
# 模型和分词器被加载到内存中，供后续所有推理请求复用。
print("="*80)
print("正在加载模型和分词器，此过程可能需要几分钟，请耐心等待...")
start_time = time.time()

# 检查模型路径是否存在，提供更友好的错误提示
if not os.path.exists(MODEL_PATH):
    raise FileNotFoundError(f"错误：模型路径 '{MODEL_PATH}' 不存在。请检查路径是否正确。")

try:
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_PATH,
        use_fast=False,
        trust_remote_code=True
    )

    # 加载模型
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        ms_dtype=mindspore.bfloat16,
        device_map=0
    )
    # 将模型移动到 NPU 上
    model = model.npu()  # type: ignore

    end_time = time.time()
    print(f"模型加载成功！耗时: {end_time - start_time:.2f} 秒。")

    # --- 模型预热 ---
    print("="*80)
    print("正在进行模型预热，首次推理编译可能需要一些时间...")
    warmup_start_time = time.time()
    try:
        # 构造一个虚拟的输入，内容不重要，但格式要对
        warmup_text = "这是一个预热测试"
        warmup_conversation = [{"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": warmup_text}]
        inputs = tokenizer.apply_chat_template(
            warmup_conversation,
            add_generation_prompt=True,
            tokenize=True,
            return_tensors="ms",
            return_dict=True
        ).to('cuda')

        # 执行一次简短的生成，触发编译
        _ = model.generate(inputs['input_ids'], max_length=64)

        warmup_end_time = time.time()
        print(f"模型预热完成！耗时: {warmup_end_time - warmup_start_time:.2f} 秒。")
    except Exception as e:
        print(f"模型预热失败: {e}")
    print("="*80)

    print("Gradio 服务即将启动...")
    print("="*80)

except Exception as e:
    print(f"加载模型时发生错误: {e}")
    # 如果模型加载失败，则退出程序
    exit()

# --- 2. 定义推理函数 ---
def predict(message: str, history: list = []):
    """
    接收用户输入的需求，使用预加载的模型进行分析并返回结果。
    """
    # 对空输入进行处理
    if not message.strip():
        return "⚠️ 输入不能为空，请输入您的需求描述。"


    prompt = f"""
    你是一位专业的知识助手，名字叫XDUMsgBot，你可以根据相关片段中的信息回答用户关于西安电子科技大学的问题。
    请根据用户的问题和下列片段生成准确的回应。

    用户问题：{message}
    
    相关片段：
    {rag_pipeline(message)}
    
    请基于上述内容作答，如果没有明确的信息可供参考，请回答不知道，不要编造信息。
    
    """
    
    # 构建聊天模板
    conversation = [{"role": "system", "content": prompt}] + [
        {"role": msg["role"], "content": msg["content"]}
        for msg in history
    ] + [{"role": "user", "content": message}]

    inputs = tokenizer.apply_chat_template(
        conversation,
        add_generation_prompt=True,
        tokenize=True,
        return_tensors="ms",
        return_dict=True,
        enable_thinking=False
    ).to('cuda')
    print(f"输入的 token 数量: {inputs['input_ids'].shape[1]}, 内容: {inputs}")

    # 定义生成参数
    gen_kwargs = {"max_length": 1024, "do_sample": True, "top_k": 1}

    # 执行推理
    with core.no_grad():
        outputs = model.generate(**inputs, **gen_kwargs)
        print(f"响应生成完毕, 共生成 {outputs.shape[1]} 个 token。")
        outputs = outputs[:, inputs['input_ids'].shape[1]:]
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print(f"模型响应: {response}")
    return response


# --- 3. 创建并启动 Gradio 界面 ---
with gr.Blocks() as demo:
    gr.Markdown(GRADIO_TITLE)
    chatbot = gr.Chatbot(
        value=[],
        elem_id="chatbot",
        label="对话窗口",
        bubble_full_width=False,
        height=600
    )

    # 将 ChatInterface 的 fn 指向我们新的流式函数 predict_streaming
    gr.ChatInterface(
        fn=predict,
        chatbot=chatbot,
        examples=GRADIO_EXAMPLES,
        title=None,
        description=None,
        textbox=gr.Textbox(placeholder="请在这里输入您的需求...", container=False, scale=7),
        submit_btn="🚀 发送",
        type="messages",
        save_history=True,
    )

demo.launch(server_name="0.0.0.0")
